"""
        测试transformer模型是否能运行
"""
#-*- coding : utf-8 -*-
# coding: utf-8
import tensorflow as tf
import jieba

import config
from transformer_model.transformer import Transformer, CustomSchedule
from data_loader import DataLoader

if __name__ == "__main__":
    # 实例化一个数据加载器
    # data_loader = DataLoader(path=config.data_path, num_examples=100000)
    data_loader = DataLoader(path=config.data_path, reverse=False)
    # 获取数据DataFrame类型
    data_frame, dataset = data_loader.load_datasets(return_dataset=True)

    # 有词典时直接导入词典
    inp_tokenizer, targ_tokenizer = data_loader.load_tokenizer(config.input_tokenizer_path,
                                                               config.target_tokenizer_path)
    input_vocab_size = inp_tokenizer.vocab_size + 2
    target_vocab_size = targ_tokenizer.vocab_size + 2
    # 将数据划分成训练集和测试集
    train, test = data_loader.split_train_test(data_frame, test_size=0.2)

    # 将数据转换为dataset
    train_dataset = data_loader.data_to_dataset(train)
    test_dataset = data_loader.data_to_dataset(test)

    # 设置学习率
    learning_rate = CustomSchedule()

    # 定义优化器
    optimizer = tf.keras.optimizers.Adam(learning_rate, beta_1=0.9, beta_2=0.98,
                                         epsilon=1e-9)
    # 定义损失函数
    loss_object = tf.keras.losses.SparseCategoricalCrossentropy(
        from_logits=True, reduction='none')
    # 实例化Transformer模型
    transformer_model = Transformer(input_vocab_size=input_vocab_size, target_vocab_size=target_vocab_size,
                                    pe_input=input_vocab_size, pe_target=target_vocab_size,
                                    optimizer=optimizer,
                                    loss_object=loss_object,
                                    learning_rate=learning_rate,
                                    input_tokenizer=inp_tokenizer,
                                    target_tokenizer=targ_tokenizer,
                                    checkpoint_path= config.test_checkpoint_path
                                    )

    # # 加载模型
    # transformer_model.load_model()
    # 训练模型
    transformer_model.train(train_dataset, 1)

    # 评估模型
    transformer_model.eval(test_dataset)
    text = "what are you doing ?"
    inp_sentence, predict_sentence = transformer_model.translate(text)
    print("输入句子:{}\n预测句子:{}".format(inp_sentence, predict_sentence))
